# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jhkim

import collections
import recordclass
import functools
import os
import datetime
import logging
import sys
import numpy
import pandas

from absl import app, flags

from concurrent.futures import ProcessPoolExecutor
from coin.exchange.okex_futures.kr_rest.futures_product import OkexFuturesProduct
from coin.exchange.huobi_futures.kr_rest.futures_product import HuobiFuturesProduct
from coin.support.proto_log.logic.util import run_from_og_log_archive
from coin.exchange.kr_rest.product.product_impl import generate_product_from_str
from coin.strategy.mm.tool.order_reader import OrderReader, OrderInfo, FillInfo
from coin.strategy.mm.fastfeature.ridge_regress import get_quantile

import coin.strategy.mm.tool.archive_base as abase
import coin.strategy.mm.feed as feed
import coin.strategy.mm.subscription as sub

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

FLAGS = flags.FLAGS

# TODO: this thing appears again and again
BboInfo = recordclass.recordclass('BboInfo',
                                  ['timestamp', 'bid0_price', 'bid0_qty', 'ask0_price', 'ask0_qty'])


def bisect_batch(lhs, rhs):
  i = 0
  j = 0
  rindex = []
  while i < len(lhs):
    while j < len(rhs):
      if lhs[i] > rhs[j]:
        j += 1
      else:
        break
    rindex.append(min(j, len(rhs) - 1))
    i += 1
  return rindex


class TakerRatioStrategy(object):
  def __init__(self, trading_date, product, nametuple):
    self._nametuple = nametuple
    self._trading_date = trading_date
    self._product = product
    # id -> (price, side, time)
    self._book_snaps = {}
    self._last_book_timestamp = 0
    self._price_info = {}
    self._order_reader = OrderReader(self._product)
    self._bbos = []

  def analyze(self):
    orders_map = self._order_reader.get_orders_map()
    rescol = []
    try:
      fdf = pandas.DataFrame(self._order_reader._fill_list, columns=FillInfo.__fields__)
      orderdf = pandas.DataFrame(list(orders_map.values()), columns=OrderInfo.__fields__)
      bbodf = pandas.DataFrame(self._bbos,
                               columns=BboInfo.__fields__).sort_values(['timestamp'
                                                                       ]).reset_index(drop=True)
    except Exception:
      fdf = pandas.DataFrame(self._order_reader._fill_list, columns=FillInfo._fields)
      orderdf = pandas.DataFrame(list(orders_map.values()), columns=OrderInfo._fields)
      bbodf = pandas.DataFrame(self._bbos,
                               columns=BboInfo._fields).sort_values(['timestamp'
                                                                    ]).reset_index(drop=True)
    orderdf = orderdf.loc[orderdf['accepted_timestamp'] > 0].reset_index(drop=True)
    bbodf['mid_price'] = 0.5 * (bbodf['ask0_price'] + bbodf['bid0_price'])

    idxproj12 = bisect_batch((orderdf['submit_timestamp'].astype(int)).tolist(),
                             bbodf['timestamp'].astype(int).tolist())
    fidxproj12 = bisect_batch((fdf['submit_timestamp'].astype(int)).tolist(),
                              bbodf['timestamp'].astype(int).tolist())
    for col in ['ask0_price', 'bid0_price', 'mid_price']:
      orderdf[col] = bbodf[col].iloc[idxproj12].reset_index(drop=True)
      fdf[col] = bbodf[col].iloc[fidxproj12].reset_index(drop=True)

    tdiff = bbodf.loc[idxproj12].reset_index(drop=True).timestamp - orderdf['submit_timestamp']
    orderdf = orderdf.loc[tdiff.abs() < 1e9].reset_index(drop=True)

    fwds = [0, 10, 50, 100, 200, 500, 1000, 2000, 3000]

    orderdf['agg_lifetime'] = 0

    prev_fwd_ms = None
    for fwd_ms in fwds + [10000, 30000, 60000]:
      fidxproj12 = bisect_batch((fdf['submit_timestamp'].astype(int) + fwd_ms * 1e6).tolist(),
                                bbodf['timestamp'].astype(int).tolist())

      idxproj12 = bisect_batch((orderdf['submit_timestamp'].astype(int) + fwd_ms * 1e6).tolist(),
                               bbodf['timestamp'].astype(int).tolist())
      for col in ['ask0_price', 'bid0_price', 'mid_price']:
        orderdf['%s_%sms' % (col, fwd_ms)] = bbodf[col].iloc[idxproj12].reset_index(drop=True)
        fdf['%s_%sms' % (col, fwd_ms)] = bbodf[col].iloc[fidxproj12].reset_index(drop=True)

      fidxbuy = (fdf['sign'] == 1)
      fidxsell = (fdf['sign'] == -1)
      fdf['buy_payup_%sms'
          % fwd_ms] = fidxbuy * (fdf['order_price'] - fdf['ask0_price_%sms' % fwd_ms])
      fdf['sell_payup_%sms'
          % fwd_ms] = fidxsell * (-fdf['order_price'] + fdf['bid0_price_%sms' % fwd_ms])

      fdf['buy_ret_%sms'
          % fwd_ms] = fidxbuy * (fdf['mid_price_%sms' % fwd_ms] - fdf['ask0_price_%sms' % fwds[0]])
      fdf['sell_ret_%sms' %
          fwd_ms] = fidxsell * (-fdf['mid_price_%sms' % fwd_ms] + fdf['bid0_price_%sms' % fwds[0]])
      fdf['ret_%sms' % fwd_ms] = fdf['buy_ret_%sms' % fwd_ms] + fdf['sell_ret_%sms' % fwd_ms]

      fdf['fill_buy_ret_%sms'
          % fwd_ms] = fidxbuy * (fdf['mid_price_%sms' % fwd_ms] - fdf['fill_price'])
      fdf['fill_sell_ret_%sms'
          % fwd_ms] = fidxsell * (-fdf['mid_price_%sms' % fwd_ms] + fdf['fill_price'])
      fdf['fill_ret_%sms'
          % fwd_ms] = fdf['fill_buy_ret_%sms' % fwd_ms] + fdf['fill_sell_ret_%sms' % fwd_ms]

      idxbuy = (orderdf['sign'] == 1)
      idxsell = (orderdf['sign'] == -1)
      orderdf['buy_payup_%sms'
              % fwd_ms] = idxbuy * (orderdf['price'] - orderdf['ask0_price_%sms' % fwd_ms])
      orderdf['sell_payup_%sms'
              % fwd_ms] = idxsell * (-orderdf['price'] + orderdf['bid0_price_%sms' % fwd_ms])

      orderdf['buy_ret_%sms' % fwd_ms] = idxbuy * (orderdf['mid_price_%sms' % fwd_ms]
                                                   - orderdf['ask0_price_%sms' % fwds[0]])
      orderdf['sell_ret_%sms' % fwd_ms] = idxsell * (-orderdf['mid_price_%sms' % fwd_ms]
                                                     + orderdf['bid0_price_%sms' % fwds[0]])

      orderdf['ret_%sms'
              % fwd_ms] = orderdf['buy_ret_%sms' % fwd_ms] + orderdf['sell_ret_%sms' % fwd_ms]

      idxfill = orderdf['filled_timestamp'] > 0
      orderdf['fill_ret_%sms'
              % fwd_ms] = (idxfill * idxbuy *
                           (orderdf['mid_price_%sms' % fwd_ms] -
                            (orderdf['fill_pq'] / orderdf['fill_qty'])) + idxfill * idxsell *
                           (-orderdf['mid_price_%sms' % fwd_ms] +
                            (orderdf['fill_pq'] / orderdf['fill_qty']))) * (1e4 / orderdf['price'])

      orderdf['buy_m2m_ret_%sms' % fwd_ms] = idxbuy * (orderdf['mid_price_%sms' % fwd_ms]
                                                       - orderdf['mid_price_%sms' % fwds[0]])
      orderdf['sell_m2m_ret_%sms' % fwd_ms] = idxsell * (-orderdf['mid_price_%sms' % fwd_ms]
                                                         + orderdf['mid_price_%sms' % fwds[0]])
      orderdf['m2m_ret_%sms' %
              fwd_ms] = orderdf['buy_m2m_ret_%sms' % fwd_ms] + orderdf['sell_m2m_ret_%sms' % fwd_ms]

      orderdf['payup_%sms'
              % fwd_ms] = orderdf['buy_payup_%sms' % fwd_ms] + orderdf['sell_payup_%sms' % fwd_ms]
      orderdf['is_agg_until_%sms' % fwd_ms] = orderdf['payup_%sms' % fwd_ms] > 0
      if prev_fwd_ms is not None:
        orderdf['is_agg_until_%sms' % fwd_ms] &= orderdf['is_agg_until_%sms' % prev_fwd_ms]
      if fwd_ms <= 2000:
        orderdf.loc[orderdf['is_agg_until_%sms' % fwd_ms], 'agg_lifetime'] = fwd_ms
      prev_fwd_ms = fwd_ms

    orderdf['implied_taker'] = ~numpy.isclose(
        (orderdf['fill_pq'] / orderdf['fill_qty']), orderdf['price']) & (orderdf['filled_timestamp']
                                                                         >= 0)

    orderdf['implied_miss'] = ~numpy.isclose(
        (orderdf['fill_pq'] / orderdf['fill_qty']), orderdf['price']) & (
            orderdf['filled_timestamp'] < 0) & ~((orderdf['price'] == orderdf['ask0_price']) |
                                                 (orderdf['price'] == orderdf['bid0_price']))

    filldf = orderdf.loc[orderdf['filled_timestamp'] > 0]
    tsdiff, absolute_taker_ratio = get_quantile(
        filldf['filled_timestamp'] - filldf['submit_timestamp'],
        ~numpy.isclose(filldf['price'], filldf['fill_pq'] / filldf['fill_qty']), 10)
    tsdiff, fill_ret_10s = get_quantile(filldf['filled_timestamp'] - filldf['submit_timestamp'],
                                        filldf['fill_ret_10000ms'], 10)

    plt.rcParams['figure.figsize'] = 14, 12
    plt.rcParams['font.family'] = ['Dejavu Sans Mono']
    plt.rcParams['font.sans-serif'] = ['Dejavu Sans Mono']
    plt.rcParams['font.monospace'] = ['Dejavu Sans Mono']
    plt.rcParams['font.size'] = 10
    plt.rcParams['legend.fontsize'] = 6
    plt.rcParams['xtick.labelsize'] = 10
    plt.rcParams['ytick.labelsize'] = 10

    plt.subplot(211)
    plt.plot(tsdiff, absolute_taker_ratio, 'g.-')
    plt.title("%s: tsdiff vs absolute taker ratio" % self._nametuple)
    plt.subplot(212)
    plt.plot(tsdiff, fill_ret_10s, 'g.-')
    plt.title("tsdiff vs 10s ret(bps)")
    plt.tight_layout()
    plt.close()

    aggprobs = []
    fwdprobs = []
    fwdtakers = []
    m2mrets = []

    passprobs = []
    nottakerprobs = []

    fdf.to_csv(f"{FLAGS.out_dir}/fills_{self._trading_date}.csv")

    for fwd_ms in fwds:
      aggprobs.append(orderdf['is_agg_until_%sms' % fwd_ms].mean())
      fwdprobs.append(
          (orderdf['filled_timestamp'].loc[orderdf['is_agg_until_%sms' % fwd_ms]] > 0).mean())
      fwdtakers.append((orderdf['implied_taker'].loc[orderdf['is_agg_until_%sms' % fwd_ms]]).mean())

      passprobs.append(1 - orderdf['is_agg_until_%sms' % fwd_ms].mean())
      nottakerprobs.append(
          1 - (orderdf['implied_taker'].loc[~orderdf['is_agg_until_%sms' % fwd_ms]]).mean())

      m2mrets.append(orderdf['m2m_ret_%sms' % fwd_ms].mean())

    plt.suptitle(self._nametuple)
    plt.subplot(321)
    plt.plot(fwds, fwdprobs, 'g.-', fwds, fwdtakers, 'g.-')
    plt.ylabel("fill ratio(red), implied taker ratio(green)")
    plt.twinx()
    plt.plot(fwds, aggprobs, 'r.-')
    plt.ylabel("ratio stays agg(blue)")
    plt.title("agg order age vs fill ratio")
    '''
    plt.subplot(322)
    plt.plot(fwds, m2mrets, 'g.-')
    plt.title("time vs m2m returns(price goes away)")
    '''
    plt.subplot(322)
    plt.plot(fwds, nottakerprobs, 'g.-')
    plt.ylabel("not (implied)taker ratio(red)")
    plt.twinx()
    plt.plot(fwds, passprobs, 'r.-')
    plt.ylabel("ratio becomes pass(blue)")
    plt.title("pass order age vs implied taker ratio")

    plt.subplot(323)
    agglifetime, takratio = get_quantile(orderdf['agg_lifetime'], orderdf['implied_taker'], 30)
    agglifetime, frqt = get_quantile(orderdf['agg_lifetime'], orderdf['filled_timestamp'] > 0, 30)
    plt.plot(agglifetime, takratio, 'g.-', agglifetime, frqt, 'g.-')
    plt.title("agg lifetime vs taker ratio & agg lifetime vs fill ratio")

    plt.subplot(324)
    among_payup_index = (orderdf['payup_0ms'] >= 0)
    porderdf = orderdf.loc[among_payup_index].reset_index(drop=True)

    takerdist = sorted(porderdf['agg_lifetime'].loc[porderdf['implied_taker']])
    missdist = sorted(porderdf['agg_lifetime'].loc[porderdf['implied_miss']])
    dists = takerdist + missdist
    if len(dists) > 0:
      labels = len(takerdist) * [1] + len(missdist) * [0]
      dists, labels = zip(*sorted(zip(dists, labels)))
      dists = pandas.Series(dists)
      labels = numpy.array(labels)
      plt.plot(dists[labels == 0], 'r.', dists[labels == 1], 'g.', markersize=1)
      plt.title("taker agg lifetime vs filled(red), missed(blue)")
    plt.subplot(325)
    payup, ret1s = get_quantile(orderdf.loc[orderdf['sign'] == -1]['sell_payup_0ms'],
                                orderdf[orderdf['sign'] == -1]['ret_1000ms'], 10)
    plt.title("sell_payup vs 1s ret")
    plt.plot(payup, ret1s, 'g.-')
    payup, m2m_ret1s = get_quantile(orderdf.loc[orderdf['sign'] == -1]['sell_payup_0ms'],
                                    orderdf[orderdf['sign'] == -1]['m2m_ret_1000ms'], 10)
    plt.plot(payup, m2m_ret1s, 'r.-')
    plt.subplot(326)
    payup, ret1s = get_quantile(orderdf.loc[orderdf['sign'] == 1]['buy_payup_0ms'],
                                orderdf[orderdf['sign'] == 1]['ret_1000ms'], 10)
    plt.title("buy__payup vs 1s ret")
    plt.plot(payup, ret1s, 'g.-')
    payup, m2m_ret1s = get_quantile(orderdf.loc[orderdf['sign'] == 1]['buy_payup_0ms'],
                                    orderdf[orderdf['sign'] == 1]['m2m_ret_1000ms'], 10)
    plt.plot(payup, m2m_ret1s, 'r.-')
    '''
    plt.subplot(325)
    payup, m2mret10s = get_quantile(orderdf['payup_0ms'], orderdf['m2m_ret_10000ms'], 10)
    plt.plot(payup, m2mret10s, 'g.-')
    plt.title("payup vs 10s m2m ret")
    plt.subplot(326)
    payup, m2mret1s = get_quantile(orderdf['payup_0ms'], orderdf['m2m_ret_1000ms'], 10)
    plt.plot(payup, m2mret1s, 'g.-')
    plt.title("payup vs 1s m2m ret")
    '''
    plt.tight_layout()
    plt.savefig(f"{FLAGS.out_dir}/tr_prob_%s.png" % self._nametuple)
    plt.close()

    fq, fret_10s = get_quantile(filldf['fill_qty'], filldf['fill_ret_10000ms'], 10)
    plt.plot(fq, fret_10s, 'g.-')
    plt.title(self._nametuple)
    plt.savefig(f"{FLAGS.out_dir}/tr_adverse_%s.png" % self._nametuple)
    plt.close()

    sign_name = {-1: 'sell', 1: 'buy'}

    for j, sign in enumerate([-1, 1]):
      idxbs = orderdf['sign'] == sign
      signdf = orderdf.loc[idxbs].reset_index(drop=True)
      i = 0
      for i, fwd_ms in enumerate(fwds):
        payup = orderdf['payup_%sms' % fwd_ms]
        filled = orderdf['filled_timestamp'] > 0
        puqtl, ftl = get_quantile(payup, filled, 10)
        crosses = payup >= 0
        crossfill = filled[crosses]
        plt.subplot(len(fwds), 2, i * 2 + j + 1)
        plt.plot(puqtl, ftl, 'g.-', lw=0.5)
        plt.title("%s payup(vs price %4dms after) vs FR(crosses %02d%% crossfill %02d%%)" %
                  (sign_name[sign],
                   fwd_ms,
                   numpy.nan_to_num(crosses.mean() * 100),
                   numpy.nan_to_num(crossfill.mean() * 100)))

    os.makedirs(FLAGS.out_dir, exist_ok=True)
    plt.tight_layout()
    plt.close()

    buy_payup_avg = orderdf.loc[orderdf['sign'] == 1]['buy_payup_0ms'].mean()
    sell_payup_avg = orderdf.loc[orderdf['sign'] == -1]['sell_payup_0ms'].mean()

    stats = []

    for col in orderdf.columns:
      if not col.startswith("fill_ret"):
        continue

      takqtys = orderdf['fill_qty'][orderdf['implied_taker'] & (orderdf['fill_qty'] > 0)]
      makqtys = orderdf['fill_qty'][~orderdf['implied_taker'] & (orderdf['fill_qty'] > 0)]

      if takqtys.sum() > 0:
        takret = numpy.average(orderdf[col][orderdf['implied_taker'] & (orderdf['fill_qty'] > 0)],
                               weights=takqtys)
      else:
        takret = 0

      if makqtys.sum() > 0:
        makret = numpy.average(orderdf[col][~orderdf['implied_taker'] & (orderdf['fill_qty'] > 0)],
                               weights=makqtys)
      else:
        makret = 0

      missgrp = orderdf.groupby('implied_miss').mean()[col.replace("fill_ret", "m2m_ret")]

      try:
        missm2mret = missgrp[True]
        missqty = orderdf.groupby('implied_miss').sum()['order_qty'][True]
      except Exception:
        missm2mret = 0
        missqty = 0
      try:
        fillm2mret = missgrp[False]
        fillqty = orderdf.groupby('implied_miss').sum()['fill_qty'][False]
      except Exception:
        fillm2mret = 0
        fillqty = 0

      # print(get_quantile(orderdf['ask0_price'] - orderdf['bid0_price'], numpy.nan_to_num(orderdf['fill_ret_30000ms']), 10))
      # import pdb; pdb.set_trace()

      stats.append(
          collections.OrderedDict([
              ('trading_date', self._trading_date),
              ('horiz', col),
              ('takret', takret),
              ('makret', makret),
              ('takqty', takqtys.sum()),
              ('makqty', makqtys.sum()),
              ('missm2mret', missm2mret),
              ('fillm2mret', fillm2mret),
              ('missqty', missqty),
              ('fillqty', fillqty),
              ('buy_payup_avg', buy_payup_avg),
              ('sell_payup_avg', sell_payup_avg),
          ]))

    df = pandas.DataFrame(stats)
    print(df.set_index(['trading_date']))
    df.to_csv(f"{FLAGS.out_dir}/fillstat_{self._trading_date}.csv")

  def on_book_reset(self, book_builder_name, book_builder):
    book_builder.subscribe(self._product, functools.partial(self.on_book, self._product))

  def on_book(self, product, book):
    if product != self._product:
      return
    if self._last_book_timestamp + 5e8 > book.timestamp:
      return
    if book.ask0() is None or book.bid0() is None:
      return
    self._last_book_timestamp = book.timestamp
    self._bbos.append(
        BboInfo(book.timestamp,
                book.bid0().price,
                book.bid0().qty,
                book.ask0().price,
                book.ask0().qty))

  def on_log(self, *args, **kwargs):
    self._order_reader.on_log2(*args, **kwargs)


def launch(baseflags, fsflags, osflags, symbol, postfix):
  product = abase.symbols_to_products([symbol], baseflags)[0]
  start_time, end_time = abase.get_time_range(baseflags)

  nametuple = "_".join((baseflags.exchange,
                        osflags.orderlog_machine,
                        symbol,
                        baseflags.trading_date,
                        osflags.owner or "",
                        postfix or ""))
  strategy = TakerRatioStrategy(baseflags.trading_date, product, nametuple)

  abase.run_from_order_archive_base(baseflags, osflags, on_log=strategy.on_log)

  if len(strategy._order_reader._orders_map) == 0:
    logging.info("%s: no order event, skip.", baseflags.trading_date)
    return

  abase.run_from_archive_base(baseflags=baseflags,
                              feed_sub_flags_proto=fsflags,
                              products=[product],
                              on_book_reset=strategy.on_book_reset,
                              api_version_override={('Spot', 'Binance'): None})

  strategy.analyze()
  return 0


def launch_date(trading_date, postfix=None):
  logging.info(trading_date)
  baseflags = abase.get_base_flags()
  baseflags.trading_date = trading_date.strftime("%Y%m%d")
  fsflags = abase.get_feed_subscription_flags()
  osflags = abase.get_order_subscription_flags()
  launch(baseflags, fsflags, osflags, flags.FLAGS.symbol, postfix)


def get_trading_dates(trading_date_str):
  if trading_date_str.find("-") >= 0:
    tdfrom, tdto = trading_date_str.split("-")
    return pandas.date_range(tdfrom, tdto).to_pydatetime()
  else:
    return [datetime.datetime.strptime(trading_date_str, '%Y%m%d')]


def main(_):
  # import warnings; warnings.filterwarnings("error")
  exchange = FLAGS.exchange
  trading_date = FLAGS.trading_date
  machine = FLAGS.orderlog_machine
  assert machine, '--orderlog_machine must be specified.'
  assert exchange, '--exchange must be specified.'
  print('Running for %s %s %s ...' % (trading_date, machine, exchange))
  sys.stdout.flush()
  trading_dates = get_trading_dates(trading_date)
  if len(trading_dates) == 1:
    launch_date(trading_dates[0])
  else:
    with ProcessPoolExecutor(max_workers=FLAGS.cpu) as executor:
      futures = []
      for trading_date in trading_dates:
        futures.append(executor.submit(launch_date, trading_date))
      for future in futures:
        future.result()


if __name__ == '__main__':
  abase.define_base_flags()
  abase.define_order_archive_flags()
  abase.define_feed_archive_flags()

  flags.FLAGS.feed_machine = 'feed-01.ap-northeast-1.aws'
  flags.DEFINE_string('symbol', 'EOS-USD.QUARTER', 'symbol name')
  flags.DEFINE_integer('cpu', 4, 'num cpu to use')
  flags.DEFINE_string('out_dir', 'oe_stat_tr', 'out dir name')

  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')
  app.run(main)
